#!/usr/bin/env python3
"""Swing-up and PID stabilisation controller driving the arm render client.

This script connects to the Arm stream server via :class:`RenderClient` and
computes torques for the actuated joint using a two-stage control scheme: a
heuristic swing-up phase followed by a configurable PID stabiliser. The
controller operates on the pendulum angle (``phi``) extracted with the same
logic as the renderer to keep the behaviour consistent with the built-in
visualisation.
"""

from __future__ import annotations

import argparse
import logging
import math
import signal
import sys
from dataclasses import dataclass
from enum import Enum
from typing import Optional, Tuple

import numpy as np

from parnassus.clients.arm_stream.render_client import RenderClient

LOGGER = logging.getLogger(__name__)


class ControlPhase(Enum):
    SWING_UP = "swing_up"
    STABILISE = "stabilise"


@dataclass
class PIDController:
    """Minimal PID controller with output saturation and integral clamping."""

    kp: float
    ki: float
    kd: float
    dt: float
    setpoint: float = 0.0
    output_limits: Tuple[float, float] = (-1.0, 1.0)

    _integral: float = 0.0
    _prev_error: float = 0.0
    _initialised: bool = False

    def reset(self, *, setpoint: Optional[float] = None) -> None:
        if setpoint is not None:
            self.setpoint = setpoint
        self._integral = 0.0
        self._prev_error = 0.0
        self._initialised = False

    def __call__(self, measurement: float, measurement_dot: Optional[float]) -> float:
        error = self.setpoint - measurement

        self._integral += error * self.dt

        if measurement_dot is None:
            derivative = 0.0 if not self._initialised else (error - self._prev_error) / self.dt
        else:
            # d/dt(setpoint - measurement) = -measurement_dot when setpoint is constant
            derivative = -measurement_dot

        output = self.kp * error + self.ki * self._integral + self.kd * derivative

        low, high = self.output_limits
        clamped = np.clip(output, low, high)
        if clamped != output and self.ki != 0.0:
            # Back-calculate integral to keep the controller from winding up
            self._integral = (clamped - self.kp * error - self.kd * derivative) / self.ki
            output = clamped
        else:
            output = clamped

        self._prev_error = error
        self._initialised = True
        return float(output)


@dataclass
class SwingUpController:
    """Heuristic swing-up controller to pump energy before balancing."""

    amplitude: float
    frequency: float
    dt: float

    def __call__(self, step: int, theta: float, theta_dot: float) -> float:
        time_s = step * self.dt

        # 基础正弦波控制
        control = self.amplitude * math.sin(2.0 * math.pi * self.frequency * time_s)

        # 增强的速度反馈项
        velocity_gain = 0.8 * self.amplitude
        control += velocity_gain * np.tanh(theta_dot)

        # 基于角度的能量注入 - 在下摆时增加更多力矩
        if abs(theta) > math.pi / 2:  # 当摆在下半部分时
            energy_boost = 0.6 * self.amplitude * np.sign(theta_dot) * (1 + abs(math.cos(theta)))
            control += energy_boost

        return float(np.clip(control, -self.amplitude, self.amplitude))


def _extract_angles(
    client: RenderClient,
    observation: Optional[np.ndarray],
    info: Optional[dict],
) -> Tuple[float, float]:
    if observation is None:
        raise ValueError("Cannot compute control without an observation from the server.")
    extractor = getattr(client, "_state_extractor")
    theta, phi = extractor(np.asarray(observation), info)
    return float(theta), float(phi)


def run_pid_episode(
    client: RenderClient,
    controller: PIDController,
    swing: SwingUpController,
    *,
    seed: Optional[int],
    max_steps: int,
    swing_threshold: float,
    swing_velocity_threshold: float,
) -> None:
    obs = client.reset(seed=seed)
    if obs is None:
        raise RuntimeError("Server returned empty observation on reset; aborting episode.")

    controller.reset()
    info = None
    phase = ControlPhase.SWING_UP

    for step in range(max_steps):
        theta, phi = _extract_angles(client, obs, info)
        theta_dot = float(obs[2]) if len(obs) >= 3 else 0.0
        phi_dot = float(obs[3]) if len(obs) >= 4 else None

        if phase is ControlPhase.SWING_UP:
            control = swing(step, theta, theta_dot)
            if phi_dot is not None:
                ready = abs(phi) < swing_threshold and abs(phi_dot) < swing_velocity_threshold
            else:
                ready = abs(phi) < swing_threshold

            if ready:
                LOGGER.info(
                    "Switching to PID stabilisation at step %d (phi=%.3f, phi_dot=%s)",
                    step,
                    phi,
                    "n/a" if phi_dot is None else f"{phi_dot:.3f}",
                )
                controller.reset()
                phase = ControlPhase.STABILISE
        else:
            control = controller(phi, phi_dot)

        action = np.array([control], dtype=np.float32)

        reply = client.step(action)
        obs = reply.get("observation")
        print(obs)
        info = reply.get("info")
        reward = reply.get("reward")
        terminated = reply.get("terminated")
        truncated = reply.get("truncated")

        LOGGER.debug(
            "step=%d phase=%s theta=%.3f phi=%.3f phi_dot=%s action=%.3f reward=%s terminated=%s truncated=%s",
            step,
            phase.value,
            theta,
            phi,
            None if phi_dot is None else f"{phi_dot:.3f}",
            control,
            reward,
            terminated,
            truncated,
        )

        if obs is None:
            LOGGER.warning("Server returned no observation at step %d; stopping episode.", step)
            break

        if terminated or truncated:
            LOGGER.info(
                "Episode ended early at step %d (terminated=%s, truncated=%s).",
                step,
                terminated,
                truncated,
            )
            break


def _parse_args(argv: Optional[list[str]]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run a PID-controlled Arm render client.")
    parser.add_argument("--server", default="localhost:50051", help="gRPC server address.")
    parser.add_argument("--episodes", type=int, default=1, help="Number of episodes to run.")
    parser.add_argument("--max-steps", type=int, default=2500, help="Maximum steps per episode.")
    parser.add_argument("--seed", type=int, default=None, help="Optional reset seed.")
    parser.add_argument("--setpoint", type=float, default=0.0, help="Target pendulum angle (rad).")
    parser.add_argument("--kp", type=float, default=50.0, help="Proportional gain.")
    parser.add_argument("--ki", type=float, default=0.0, help="Integral gain.")
    parser.add_argument("--kd", type=float, default=5.5, help="Derivative gain.")
    parser.add_argument("--dt", type=float, default=0.01, help="Controller update period (s).")
    parser.add_argument(
        "--action-limit",
        type=float,
        default=1.0,
        help="Absolute bound for the action sent to the server (before torque scaling).",
    )
    parser.add_argument(
        "--swing-amplitude",
        type=float,
        default=1.5,  # 从0.8增加到1.5
        help="Open-loop swing-up action amplitude (|action| limit during swing-up).",
    )
    parser.add_argument(
        "--swing-frequency",
        type=float,
        default=1.2,  # 从1增加到1.2
        help="Swing-up drive frequency in Hz.",
    )
    parser.add_argument(
        "--swing-threshold",
        type=float,
        default=0.4,  # 从0.25放宽到0.4
        help="Abs(phi) threshold (rad) to hand over control to the PID balancer.",
    )
    parser.add_argument(
        "--swing-velocity-threshold",
        type=float,
        default=2.0,  # 从1.0放宽到2.0
        help="Abs(phi_dot) threshold (rad/s) to hand over control to the PID balancer.",
    )
    parser.add_argument(
        "--log-level",
        default="info",
        choices=["debug", "info", "warning", "error", "critical"],
        help="Logging level.",
    )
    parser.add_argument(
        "--disable-render",
        action="store_true",
        help="Skip automatic rendering updates (useful when running headless).",
    )
    return parser.parse_args(argv)


def main(argv: Optional[list[str]] = None) -> int:
    args = _parse_args(argv)
    logging.basicConfig(
        level=getattr(logging, args.log_level.upper()),
        format="[%(levelname)s] %(asctime)s %(name)s: %(message)s",
    )

    controller = PIDController(
        kp=args.kp,
        ki=args.ki,
        kd=args.kd,
        dt=args.dt,
        setpoint=args.setpoint,
        output_limits=(-abs(args.action_limit), abs(args.action_limit)),
    )

    swing_amplitude = min(abs(args.swing_amplitude), abs(args.action_limit))
    swing = SwingUpController(
        amplitude=swing_amplitude if swing_amplitude > 0 else 0.0,
        frequency=max(0.05, args.swing_frequency),
        dt=args.dt,
    )

    LOGGER.info(
        "Starting PID control: server=%s episodes=%d max_steps=%d setpoint=%.3f",
        args.server,
        args.episodes,
        args.max_steps,
        args.setpoint,
    )
    LOGGER.debug(
        "Swing-up parameters: amplitude=%.3f frequency=%.3f threshold=%.3f velocity_threshold=%.3f",
        swing.amplitude,
        swing.frequency,
        args.swing_threshold,
        args.swing_velocity_threshold,
    )

    interrupted = False

    def _handle_sigint(signum, frame):  # type: ignore[unused-argument]
        nonlocal interrupted
        interrupted = True
        LOGGER.info("Received interrupt signal; finishing current step before exiting.")

    signal.signal(signal.SIGINT, _handle_sigint)

    auto_render = not args.disable_render

    try:
        with RenderClient(args.server, auto_render=auto_render) as client:
            for episode in range(args.episodes):
                if interrupted:
                    break
                LOGGER.info("Episode %d/%d", episode + 1, args.episodes)
                episode_seed = None if args.seed is None else int(args.seed) + episode
                run_pid_episode(
                    client,
                    controller,
                    swing,
                    seed=episode_seed,
                    max_steps=args.max_steps,
                    swing_threshold=max(0.0, args.swing_threshold),
                    swing_velocity_threshold=max(0.0, args.swing_velocity_threshold),
                )
    except KeyboardInterrupt:
        LOGGER.info("Interrupted by user.")
    except Exception as exc:  # pragma: no cover - defensive logging
        LOGGER.error("PID control loop failed: %s", exc)
        return 1

    LOGGER.info("PID run finished%s", " (interrupted)" if interrupted else "")
    return 0


if __name__ == "__main__":
    sys.exit(main())
